#pragma once

#include <memory>

#include <google/protobuf/arena.h>
#include <google/protobuf/repeated_field.h>

// copy from hawking util to avoid submodule
namespace hawking {
namespace indexlib {

class ArenaWrapper {
public:
    ArenaWrapper(std::shared_ptr<google::protobuf::Arena> arena);
    ArenaWrapper();
    ArenaWrapper(const google::protobuf::ArenaOptions& options);

    ~ArenaWrapper();

    // Can only create protocol buffer messages. If its proto file was compiled
    // with option 'cc_enable_arenas=true', the protobuf message will be created
    // on arena. Otherwise, it is an ordinary allocation followed by arena->Own(message).
    template<typename T>
    T* CreateMessageOnArena() {
        return T::default_instance().New(arena_.get());
    }

    // Its semantics are identical to CreateMessageOnArena. This will create a
    // protocol buffer RepeatedPtrMessages.
    template<typename T>
    google::protobuf::RepeatedPtrField<T>* CreateRepeatedPtrMessageOnArena() {
        return google::protobuf::Arena::Create<
            google::protobuf::RepeatedPtrField<T>>(arena_.get(), arena_.get());
    }

    // Its semantics are identical to CreateMessageOnArena. This will create a
    // protocol buffer RepeatedMessages.
    template<typename T>
    google::protobuf::RepeatedField<T>* CreateRepeatedMessageOnArena() {
        return google::protobuf::Arena::Create<
            google::protobuf::RepeatedField<T>>(arena_.get(), arena_.get());
    }

    // Create protocol buffer messages and arbitrary c++ class on arena.
    // NOTE: The underlying memory of objects will not be allocated on arena.
    template<typename T>
    T* CreateOnArena() {
        return google::protobuf::Arena::Create<T>(arena_.get());
    }

    // T MUST have a trivial constructor: constructors are not called when the
    // array is created on the arena.
    template<typename T>
    T* CreateArrayOnArena(size_t n) {
        return google::protobuf::Arena::CreateArray<T>(arena_.get(), n);
    }

    std::shared_ptr<google::protobuf::Arena> arena() { return arena_; }

    static std::shared_ptr<google::protobuf::Arena> CreateArena();
    static void DestroyArena(google::protobuf::Arena* arena);
    static void CloseMonitor();

private:
    std::shared_ptr<google::protobuf::Arena> arena_;
};


void CreateThreadLocalExclusiveArena(std::unique_ptr<indexlib::ArenaWrapper>* arena,
                                     char** arena_init_block);
void DestryThreadLocalExclusiveArena(std::unique_ptr<indexlib::ArenaWrapper>& arena,
                                     char* arena_init_block);

template<class T>
class ThreadLocalExclusivePbMessage {
public:
    ThreadLocalExclusivePbMessage() {
        CreateThreadLocalExclusiveArena(&arena_, &arena_init_block_);
        pb_msg_ = arena_->CreateMessageOnArena<T>();
    }

    ~ThreadLocalExclusivePbMessage() {
        DestryThreadLocalExclusiveArena(arena_, arena_init_block_);
        pb_msg_ = nullptr;
    }

    T* pb_msg() const {
        return pb_msg_;
    }

private:
    char* arena_init_block_ = nullptr;
    std::unique_ptr<indexlib::ArenaWrapper> arena_;
    T* pb_msg_ = nullptr;
};

} // namespace indexlib
} // namespace hawking